"""
ai_models.py
OpenAI API 的实现和配置管理
"""

import json
import os
import logging
from openai import OpenAI
from PyQt6.QtCore import QObject, pyqtSignal

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class OpenAIChat(QObject):
    """OpenAI API 聊天实现"""
    model_changed = pyqtSignal(str)  # 添加信号

    def __init__(self, config_file='model_config.json'):
        super().__init__()
        self.config_file = config_file
        self.models = self.load_config()
        self.clients = {}  # 缓存已创建的客户端
        self.current_model = None

    def load_config(self):
        """加载模型配置"""
        if os.path.exists(self.config_file):
            try:
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                logger.error(f"加载配置时发生错误: {str(e)}")
                return {}
        return {}

    def save_config(self):
        """保存模型配置"""
        with open(self.config_file, 'w', encoding='utf-8') as f:
            json.dump(self.models, f, ensure_ascii=False, indent=4)

    def add_model(self, name, api_key, base_url):
        """添加模型配置"""
        self.models[name] = {
            'api_key': api_key,
            'base_url': base_url,
            'model_name': name
        }
        self.save_config()

    def delete_model(self, name):
        """删除模型配置"""
        if name in self.models:
            del self.models[name]
            if name in self.clients:
                del self.clients[name]
            self.save_config()

    def get_model(self, name):
        """获取模型配置"""
        return self.models.get(name)

    def get_all_models(self):
        """获取所有模型配置"""
        return self.models

    def generate_response(self, model_name, prompt, history=None):
        """
        生成回复
        Args:
            model_name: 模型配置名称
            prompt: 用户输入
            history: 历史消息列表，格式为[{"role": "user", "content": "xxx"}, {"role": "assistant", "content": "xxx"}]
        Returns:
            OpenAI的流式响应
        """
        try:
            # 获取或创建客户端
            if model_name not in self.clients:
                config = self.models[model_name]
                self.clients[model_name] = OpenAI(
                    api_key=config['api_key'],
                    base_url=config['base_url']
                )

            # 构建消息列表
            messages = []
            if history:
                messages.extend(history)
            messages.append({'role': 'user', 'content': prompt})

            # 生成响应
            response = self.clients[model_name].chat.completions.create(
                model=self.models[model_name]['model_name'],
                messages=messages,
                stream=True
            )
            return response

        except Exception as e:
            raise Exception(f"OpenAI API 错误: {str(e)}")

    def create_response(self, text="", error=None, is_final=False):
        """
        创建响应对象
        Args:
            text: 响应文本
            error: 错误信息
            is_final: 是否为最终响应
        Returns:
            OpenAIResponse: 响应对象
        """
        return OpenAIResponse(text, error, is_final)

    def format_final_response(self, text):
        """
        处理最终响应的渲染
        Args:
            text: 原始响应文本
        Returns:
            str: 渲染后的文本
        """
        # 如果是代码示例的响应
        if "```" in text:
            import re
            # 匹配代码块的正则表达式，可以捕获语言标识
            pattern = r'```(\w*)\n((?:(?!```).|\n)+?)```'
            
            def replace_code_block(match):
                language = match.group(1) or 'text'  # 如果没有指定语言，默认为 text
                code = match.group(2)
                
                return f'''<div style="background-color: #f6f8fa; border-radius: 6px; margin: 10px 0; border: 1px solid #e1e4e8; overflow: hidden;">
                    <div style="background-color: #f1f1f1; padding: 8px 16px; border-bottom: 1px solid #e1e4e8;">{language}</div>
                    <pre style="margin: 0; padding: 16px; overflow-x: auto;"><code>{code}</code></pre>
                </div>'''
            
            # 替换所有代码块
            formatted_text = re.sub(pattern, replace_code_block, text, flags=re.DOTALL)
            return formatted_text
            
        # 如果是普通文本响应，保持原样
        return text

    def set_current_model(self, model_name):
        """设置当前模型"""
        if model_name in self.models:
            self.current_model = model_name
            self.model_changed.emit(model_name)
            return True
        return False

class OpenAIResponse:
    """OpenAI 响应数据结构"""
    def __init__(self, text="", error=None, is_final=False):
        self.text = text          # 响应文本
        self.error = error        # 错误信息
        self.is_final = is_final  # 是否为最终响应
        
    def set_complete(self):
        """标记响应已完成"""
        self.is_final = True
        
    def set_error(self, error):
        """设置错误信息"""
        self.error = error
        
    def has_error(self):
        """检查是否存在错误"""
        return self.error is not None

